import wfdb
import numpy as np
import csv
from myutils.NcUtils import fliter
from myutils.dataTypeUtils import ECG_R_list, ECGSignal, AAMI, AAMI_MIT

# TODO： 病人内模式，获取MLII导联所有数据
# 所有数据
# file_name = ['100', '101', '102', '103', '104', '105', '106', '107',
#              '108', '109', '111', '112', '113', '114', '115', '116',
#              '117', '118', '119', '121', '122', '123', '124', '200',
#              '201', '202', '203', '205', '207', '208', '209', '210',
#              '212', '213', '214', '215', '217', '219', '220', '221',
#              '222', '223', '228', '230', '231', '232', '233', '234']

# 北交论文中的病人间分类
# 训练集
# file_name = ['101', '106', '108', '109', '112', '114', '115', '116',
#              '118', '119', '122', '124', '201', '203', '205', '207',
#              '208', '209', '215', '220', '223', '230']
# 测试集
file_name = ['100', '103', '105', '111', '113', '117', '121', '123',
             '200', '202', '210', '212', '213', '214', '219', '221',
             '222', '228', '231', '232', '233', '234']

for F_name in file_name:
    print(f'正在处理数据{F_name}')
    signal_annotation = wfdb.rdann(f'../data/mit-bih-arrhythmia-database-1.0.0/{F_name}', "atr", sampfrom=0,
                                   sampto=650000)
    record = wfdb.rdrecord(f'../data/mit-bih-arrhythmia-database-1.0.0/{F_name}',
                           sampfrom=0, sampto=650000,
                           physical=True,
                           channels=[0, 1])

    ecg = record.p_signal.T[0]  # 现在形状是 (2, 32)
    # 五点平滑滤波
    ecg = fliter(ecg)
    # 获取表示R点的心拍类型的索引
    Index = np.isin(signal_annotation.symbol, ECG_R_list)
    # 将标签从list列表类型转化为数组类型为了用下面的索引值提取
    Label = np.array(signal_annotation.symbol)
    # 提取表示为R点的心拍标签
    Label = Label[Index]
    # 提取表示为R点的坐标
    Sample = signal_annotation.sample[Index]
    # 获取心拍种类
    Label_kind = list(set(Label))
    # 读取每一种R点在信号中的位置
    length = len(record.p_signal)
    print(f'{F_name}有效心拍总长度 = {len(Sample)}')
    for k in Label_kind:
        index = [i for i, x in enumerate(Label) if x == k]
        Signal_index = [Sample[i] for i in index]
        j = 0
        for site in Signal_index:
            if 100 < site < length - 168:
                ECGSignal[str(k)].append(ecg.tolist()[site - 100:site + 168])
            else:
                print(site)
                print(Label[index[j]])
            j = j + 1

# 打印种类
for key, value in ECGSignal.items():
    print(f'{key} = {len(value)}')

for ECG_key, ECG_value in ECGSignal.items():
    for AAMI_MIT_key, AAMI_MIT_value in AAMI_MIT.items():
        if ECG_key in AAMI_MIT_value:
            AAMI[AAMI_MIT_key].extend(ECG_value)

# 5分类并保存到CSV格式的文件里
for key, value in AAMI.items():
    # /data/mit/alldata目录需提前创建，csv自动生成
    with open(f'../data/mit/brn268/{key}_test.csv', 'w', newline='\n') as f:
        writer = csv.writer(f)
        writer.writerows(value)
